﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.WindowsAzure.StorageClient;
using System.Linq.Expressions;
using System.Data.Services.Client;
using System.Threading;

namespace AntsCode.StorageClientExtensions.Helpers
{
    internal class QueryExecutor
    {
        public QueryExecutor(EventHandler<XTableServiceSendingRequestEventArgs> sendingRequest, RetryPolicy retryPolicy)
        {
            this.SendingRequest = sendingRequest;
            this.RetryPolicy = retryPolicy;
        }

        public IQueryable<T> ExecuteAll<T>(DataServiceQuery<T> query, ref ContinuationToken continuationToken)
        {
            // Initialise the responses list with an initial capacity of 5
            List<List<T>> responses = new List<List<T>>(5);
            int maxRecords = -1; // All records
            int recordCount = 0;
            int currentBatchNo = this.QueryBatchNo;

            // Does this expression contain a Take method?
            var take = ExpressionUtil.FindMethodCall(query.Expression, new string[] { "Take" });
            ConstantExpression takeValue;

            if (take != null)
            {
                // Yes, to get how many records should be returned
                takeValue = (ConstantExpression)take.Arguments[1];
                maxRecords = (int)takeValue.Value;
            }

            // Start table service queries
            DataServiceQuery<T> currentQuery;
            QueryOperationResponse<T> response = null;
            List<T> results;

            do
            {
                currentQuery = query;

                // Add any continuation token
                if (continuationToken != null)
                {
                    currentQuery = this.AddContinuationToken(currentQuery, continuationToken);
                }

                // Fire SendingRequest event
                if (this.SendingRequest != null)
                {
                    this.SendingRequest(this, new XTableServiceSendingRequestEventArgs()
                    {
                        AsyncOrder = currentBatchNo,
                        RequestUri = currentQuery.RequestUri
                    });
                }

                // Execute the query using the supplied retry policy
                response = ExecuteWithRetries<T>(currentQuery, this.RetryPolicy);

                // Add the response to the responses list
                results = response.ToList();

                responses.Add(results);

                // Update continuation token
                continuationToken = ContinuationToken.GetFromQueryResponse(response);

                // Update record count
                recordCount += results.Count();

                if (maxRecords != -1 && recordCount < maxRecords)
                {
                    // Update the query to only return the remaining required records
                    takeValue = Expression.Constant(maxRecords - recordCount);
                }
            }
            while ((maxRecords == -1 || recordCount < maxRecords) && continuationToken != null);

            // Concatenate the results
            List<T> records = new List<T>(recordCount);

            foreach (List<T> result in responses)
            {
                if (result == null)
                {
                    // We've reached the end of the results, so don't bother iterating through null items
                    break;
                }

                records.AddRange(result);
            }

            return records.AsQueryable<T>();
        }

        private QueryOperationResponse<T> ExecuteWithRetries<T>(DataServiceQuery<T> query, RetryPolicy retryPolicy)
        {
            QueryOperationResponse<T> response = null;
            ShouldRetry shouldRetry = retryPolicy();
            TimeSpan delay;
            bool success = false;
            int retryCount = 0;

            while (!success)
            {
                try
                {
                    response = (QueryOperationResponse<T>)query.Execute();
                    success = true;
                }
                catch (Exception e)
                {
                    // Call the shouldRetry delegate to see if the query should be retried
                    if (shouldRetry(retryCount, e, out delay))
                    {
                        // Retry the query, but wait for the specified delay first
                        Thread.Sleep(delay);

                        // Increment the retry count
                        retryCount++;
                    }
                    else
                    {
                        // Do not retry, throw exception
                        throw e;
                    }
                }
            }

            return response;
        }

        private DataServiceQuery<T> AddContinuationToken<T>(DataServiceQuery<T> data, ContinuationToken continuationToken)
        {
            // Has a continuation token been provided?
            if (continuationToken != null)
            {
                // Yes, so add the relevant query options
                if (continuationToken.NextPartitionKey != null)
                {
                    data = data.AddQueryOption("NextPartitionKey", continuationToken.NextPartitionKey);
                }

                if (continuationToken.NextRowKey != null)
                {
                    data = data.AddQueryOption("NextRowKey", continuationToken.NextRowKey);
                }
            }

            return data;
        }

        public void IncrementQueryBatchNo()
        {
            this.QueryBatchNo++;
        }

        public RetryPolicy RetryPolicy
        {
            get;
            private set;
        }

        public EventHandler<XTableServiceSendingRequestEventArgs> SendingRequest;

        public int QueryBatchNo
        {
            get;
            private set;
        }
    }
}